import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans

# 生成虚拟数据集
np.random.seed(42)
X, _ = make_blobs(n_samples=3000, centers=3, cluster_std=1.0, random_state=42)

# K-Means 聚类
kmeans = KMeans(n_clusters=3, init='random', max_iter=300, n_init=10, random_state=42)
kmeans.fit(X)

# 获取聚类结果
labels = kmeans.labels_
centers = kmeans.cluster_centers_
inertia = kmeans.inertia_

# WCSS 收敛曲线
inertia_values = []
kmeans_temp = KMeans(n_clusters=3, init='random', max_iter=1, n_init=1, random_state=42)

for i in range(1, 20):
    kmeans_temp.max_iter = i
    kmeans_temp.fit(X)
    inertia_values.append(kmeans_temp.inertia_)

# 绘制图像
plt.figure(figsize=(14, 6))

# 子图 1：聚类结果
plt.subplot(1, 2, 1)
for cluster_idx in range(3):
    cluster_points = X[labels == cluster_idx]
    plt.scatter(cluster_points[:, 0], cluster_points[:, 1], label=f'Cluster {cluster_idx}', s=50)
plt.scatter(centers[:, 0], centers[:, 1], c='black', marker='X', s=200, label='Centers')
plt.title('K-Means Clustering Results', fontsize=14)
plt.xlabel('Feature 1')
plt.ylabel('Feature 2')
plt.legend()

# 子图 2：WCSS 收敛曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, 20), inertia_values, marker='o', color='red', linewidth=2)
plt.title('WCSS Convergence Curve', fontsize=14)
plt.xlabel('Iterations')
plt.ylabel('WCSS')

plt.tight_layout()
plt.show()